# pip install pika==1.3.1
import os
import sys

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, BASE_DIR)

from tools.auth_task import auth_task_state
from tools.utils import save_result_byte_file, save_result_list, save_result_list_to_server
import requests
from generate.process_generate import process_generate
import json
import settings
import traceback
import threading
import random
import base64
import time
from tools.log import logger

import copy
import pika
from pika.adapters.blocking_connection import BlockingChannel
from pika.exceptions import StreamLostError, ConnectionClosedByBroker, AMQPConnectionError, ChannelWrongStateError


# 优先级队列10
queue_arguments = {"x-max-priority": 10}

def init_rabbit_serve():
    
    # 1 链接rabbitmq
    credentials = pika.PlainCredentials(settings.QUEUE_USER, settings.QUEUE_PWD)
    connection = pika.BlockingConnection(pika.ConnectionParameters(settings.QUEUE_HOST, port=settings.QUEUE_PORT, 
                                                                   credentials=credentials, heartbeat=settings.HEARTBEAT))
    channel = connection.channel()



    for queue in settings.LISTEN_QUEUE:
        # 2 创建队列  如果已经有这个队列就不执行  如果生产者没有这个队列就创建一个队列
        channel.queue_declare(queue=queue, arguments=queue_arguments, durable=True)


    # 设置为空闲取消息
    channel.basic_qos(prefetch_count=1)
    
    return channel, connection


def new_ch_queue_send_message(data: dict, queue_name: str = settings.RESULT_QUEUE_NAME) -> None:
    connection = None
    try:
        channel, connection = init_rabbit_serve()
        print("send message")
        channel.queue_declare(queue=queue_name, arguments=queue_arguments, durable=True)
        channel.basic_publish(exchange='',
                        routing_key=queue_name,  # 指定队列
                        body=json.dumps(data),
                        properties=pika.BasicProperties(delivery_mode=2, priority=1) # 设置持久化缓存和优先级为1
                        )
    except Exception as e:
        print("send message error")
        logger.error("send message error")
        logger.error(traceback.format_exc())
    finally:
        if connection:
            connection.close()


def queue_send_message(ch: BlockingChannel, data: dict, queue_name: str = settings.RESULT_QUEUE_NAME) -> None:
    try:
        print("send message")
        ch.queue_declare(queue=queue_name, arguments=queue_arguments, durable=True)
        ch.basic_publish(exchange='',
                        routing_key=queue_name,  # 指定队列
                        body=json.dumps(data),
                        properties=pika.BasicProperties(delivery_mode=2, priority=1) # 设置持久化缓存和优先级为1
                        )
    except Exception as e:
        new_ch_queue_send_message(data, queue_name)
        # print("send message error")
        # logger.error("send message error")
        # logger.error(traceback.format_exc())


def new_channel_basic_ack(delivery_tag):
    connection = None
    try:
        print('new channel send message')
        channel, connection = init_rabbit_serve()
        channel.basic_ack(delivery_tag=delivery_tag)
    except Exception as e:
        print("send message error")
        logger.error("send message error")
        logger.error(traceback.format_exc())
    finally:
        if connection:
            connection.close()


# 定义回调函数
def callback_function(ch, method, properties, body):
    delivery_tag = method.delivery_tag
    try:
        print("收到消息")
        # 具体操作....
        data = json.loads(body.decode("utf-8"))
        generate_function(ch, data)
    except Exception as e:
        print(traceback.format_exc())
        
    finally:
        ch.basic_ack(delivery_tag=delivery_tag)
        

def start_serve():
    
    connection = None
    channel = None
    
    while True:
        
        try:
            channel, connection = init_rabbit_serve()
        
            for queue in settings.LISTEN_QUEUE:
                # 监听配置  有数据后执行回调
                channel.basic_consume(queue=queue,
                                    on_message_callback=callback_function)  # 回调函数
            print("rabbitMQ start", settings.QUEUE_HOST, settings.QUEUE_PORT)
            print("listen queue: ", settings.LISTEN_QUEUE)
            logger.info("rabbitMQ start")
            # 启动监听
            channel.start_consuming()
        except StreamLostError:
            print(traceback.format_exc())
            print("rabbitMQ Connection lost, reconnecting...")
            time.sleep(5)  # 等待后重试
            continue
        except ChannelWrongStateError:
            print(traceback.format_exc())
            print("chanel error, reconnecting...")
            time.sleep(10)  # 等待后重试
            continue
        except AMQPConnectionError:
            print(traceback.format_exc())
            print("rabbitMQ close, reconnecting...")
            time.sleep(10)  # 等待后重试
            continue
        except KeyboardInterrupt:
            break
        except Exception as e:
            print(traceback.format_exc())
            print(f"Unexpected error: {e}")
            break
        finally:
            if channel and channel.is_open:
                channel.close()
            if connection and not connection.is_closed:
                connection.close()
                print("Connection closed.")
    if channel and channel.is_open:
            channel.close()
    if connection and not connection.is_closed:
        channel.close()
        connection.close()
        print("Connection closed.")



def send_progress_message(ch: BlockingChannel, message) -> None:
    progress_dict = copy.copy(message)
    progress_dict['progress'] = 'start_generate'
    progress_ret_dict = progress_dict
    queue_send_message(ch, progress_ret_dict, queue_name=settings.PROGRESS_QUEUE_NAME)
# # ----------------------------------- 生成函数 --------------------------------------------------------


def generate_function(ch: BlockingChannel, message_data: dict, *args, **kwargs) -> None:

    logger.info(f"请求参数：{message_data}")
    print(message_data)
    ret_dict = dict()
    ret_dict["message"] = "未处理"
    ret_dict["result"] = None
    ret_dict['resp'] = message_data.get('resp', {})

    task_id = message_data.get('task_id')
    ret_dict["task_id"] = task_id

    # 验证任务状态
    if auth_task_state(task_id):
        print('验证任务失败')
        return

    # time.sleep(15)
    send_progress_message(ch, ret_dict)
    try:
        process_user_data_image_to_byte(message_data)

        result_data = process_generate(message_data)
        result_list = result_data.get('result')

        # time.sleep(40)
        # path = r'C:\Users\mojia\Desktop\avatar.png'
        # with open(path, 'rb') as f:
        #     data_b = f.read()
        # result_list = [{'data': data_b, 'type': 'png'}]

        if result_list:
            save_result_list(result_list, settings.SAVE_BASE_PATH)

            send_result_list = save_result_list_to_server(result_list)
            ret_dict['result'] = send_result_list
            ret_dict["message"] = "成功"
            new_ch_queue_send_message(ret_dict)
            return
        else:
            ret_dict["message"] = "系统生成失败！请稍后再试!"

            new_ch_queue_send_message(ret_dict)
            return
    except Exception as e:
        print(traceback.format_exc())
        logger.error(f"{traceback.format_exc()}")
        ret_dict['result'] = None
        ret_dict["message"] = "系统生成失败！请稍后再试!"

        new_ch_queue_send_message(ret_dict)


def process_user_data_image_to_byte(message_data):

    user_data = message_data.get('user_data', [])

    ret_li = []
    for i in range(len(user_data)):
        node = user_data[i]
        tmp_dict = dict()
        tmp_dict['nodeId'] = node['nodeId']
        tmp_dict['fieldName'] = node['fieldName']
        if node['fieldName'] in ['image', 'video', 'audio']:
            tmp_dict['fieldValue'] = requests.get(node['fieldValue']).content
        else:
            tmp_dict['fieldValue'] = node['fieldValue']

        ret_li.append(tmp_dict)
    message_data['user_data'] = ret_li


